import cv2
import numpy as np
import ctypes
from ctypes import *
lib = ctypes.cdll.LoadLibrary("./mnnsoftmax.so")
def CSoftmax(x):
    shape=x.shape
    size=1
    for s in shape:
        size*=s
    dim=shape[-1]
    x = x.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
    out=np.zeros(shape,dtype = np.float32)
    out = out.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
    lib.softmax(x,out, size, dim)
    out=np.ctypeslib.as_array(out, shape=shape)
    return out
        
def softmax(x):
    tmp=np.max(x,axis=-1)
    tmp=np.expand_dims(tmp,axis=-1)
    y=x-tmp
    y=np.exp(y)
    tmp=np.sum(y,axis=-1)
    tmp=np.expand_dims(tmp,axis=-1)
    y/=tmp
    return y


def testerr(x):
    out2=CSoftmax(x)
    out1=softmax(x)
    maxdist=np.max(abs(out1-out2))
    meandist=np.mean(abs(out1-out2))
    print("max dist:",maxdist)
    print("mean dist:",meandist)

max_v=0.
mean_v=0.
size=2000
max_x=np.random.uniform(-1,1,size=(1,size))
for i in range(100):
    x=np.random.uniform(-1,1,size=(1,size))
    x=x.astype(np.float32)
    out2=CSoftmax(x)
    out1=softmax(x)
    maxdist=np.max(abs(out1-out2))
    meandist=np.mean(abs(out1-out2))
    if maxdist>max_v:
       max_v=maxdist
       mean_v=meandist
       max_x=x
print("max dist:",max_v)
print("mean dist:",mean_v)
name="(-1,1)2000.npy"
np.save(name,max_x)
print("test:")
testerr(np.load(name))
